import asyncio
import json
import ssl
from websockets.server import serve

# 存储连接的客户端
clients = {}

async def handle_client(websocket):
    try:
        # 接收客户端ID
        client_id = await websocket.recv()
        clients[client_id] = websocket
        print(f"客户端 {client_id} 已连接")
        
        # 发送在线客户端列表
        peer_list = [pid for pid in clients.keys() if pid != client_id]
        await websocket.send(json.dumps({
            "type": "peers",
            "peers": peer_list
        }))

        # 接收并转发消息
        async for message in websocket:
            data = json.loads(message)
            
            # 处理不同类型的消息
            if data.get("target"):
                if data["target"] in clients:
                    # 转发消息给指定客户端
                    await clients[data["target"]].send(message)
                    print(f"消息从 {client_id} 转发到 {data['target']}")
            
    except Exception as e:
        print(f"错误: {e}")
    finally:
        if client_id in clients:
            del clients[client_id]
            # 通知其他客户端
            for other in clients.values():
                await other.send(json.dumps({
                    "type": "peer_disconnected",
                    "peer": client_id
                }))
        print(f"客户端 {client_id} 已断开连接")

async def main():
    ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ssl_context.load_cert_chain('certificate/server.crt', 'certificate/server.key')
    
    async with serve(handle_client, "0.0.0.0", 8765, ssl=ssl_context):
        print("WebSocket信令服务器已启动，监听端口 8765")
        await asyncio.Future()  # 保持服务器运行

if __name__ == "__main__":
    asyncio.run(main())